{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "56341c66",
   "metadata": {},
   "source": [
    "# Chapter 1: Toy Language and AST\n",
    "\n",
    "This is an xDSL version of the Toy compiler, as described in the \n",
    "[MLIR tutorial](https://mlir.llvm.org/docs/Tutorials/Toy/). This, and the following\n",
    "series of notebooks are taken close to word-for-word verbatim from the MLIR tutorials,\n",
    "as the xDSL project mirrors the MLIR structure very closely. We hope that by using these\n",
    "tutorials you will get a better idea of both now to use xDSL, and how MLIR works.\n",
    "\n",
    "## The Language\n",
    "\n",
    "This tutorial will be illustrated with a toy language that we’ll call “Toy”\n",
    "(naming is hard...). Toy is a tensor-based language that allows you to define\n",
    "functions, perform some math computation, and print results.\n",
    "\n",
    "Given that we want to keep things simple, the codegen will be limited to tensors\n",
    "of rank <= 2, and the only datatype in Toy is a 64-bit floating point type (aka\n",
    "‘double’ in C parlance). As such, all values are implicitly double precision,\n",
    "`Values` are immutable (i.e. every operation returns a newly allocated value),\n",
    "and deallocation is automatically managed. But enough with the long description;\n",
    "nothing is better than walking through an example to get a better understanding:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e07ae44f",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_0 = \"\"\"\n",
    "def main() {\n",
    "  # Define a variable `a` with shape <2, 3>, initialized with the literal value.\n",
    "  # The shape is inferred from the supplied literal.\n",
    "  var a = [[1, 2, 3], [4, 5, 6]];\n",
    "\n",
    "  # b is identical to a, the literal tensor is implicitly reshaped: defining new\n",
    "  # variables is the way to reshape tensors (element count must match).\n",
    "  var b<2, 3> = [1, 2, 3, 4, 5, 6];\n",
    "\n",
    "  # transpose() and print() are the only builtin, the following will transpose\n",
    "  # a and b and perform an element-wise multiplication before printing the result.\n",
    "  print(transpose(a) * transpose(b));\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9e14dd76",
   "metadata": {},
   "source": [
    "Type checking is statically performed through type inference; the language only\n",
    "requires type declarations to specify tensor shapes when needed. Functions are\n",
    "generic: their parameters are unranked (in other words, we know these are\n",
    "tensors, but we don't know their dimensions). They are specialized for every\n",
    "newly discovered signature at call sites. Let's revisit the previous example by\n",
    "adding a user-defined function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "853ca52d",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_1 = \"\"\"\n",
    "# User defined generic function that operates on unknown shaped arguments.\n",
    "def multiply_transpose(a, b) {\n",
    "  return transpose(a) * transpose(b);\n",
    "}\n",
    "\n",
    "def main() {\n",
    "  # Define a variable `a` with shape <2, 3>, initialized with the literal value.\n",
    "  var a = [[1, 2, 3], [4, 5, 6]];\n",
    "  var b<2, 3> = [1, 2, 3, 4, 5, 6];\n",
    "\n",
    "  # This call will specialize `multiply_transpose` with <2, 3> for both\n",
    "  # arguments and deduce a return type of <3, 2> in initialization of `c`.\n",
    "  var c = multiply_transpose(a, b);\n",
    "\n",
    "  # A second call to `multiply_transpose` with <2, 3> for both arguments will\n",
    "  # reuse the previously specialized and inferred version and return <3, 2>.\n",
    "  var d = multiply_transpose(b, a);\n",
    "\n",
    "  # A new call with <3, 2> (instead of <2, 3>) for both dimensions will\n",
    "  # trigger another specialization of `multiply_transpose`.\n",
    "  var e = multiply_transpose(b, c);\n",
    "\n",
    "  # Finally, calling into `multiply_transpose` with incompatible shape will\n",
    "  # trigger a shape inference error.\n",
    "  var f = multiply_transpose(transpose(a), c);\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7736a3ab",
   "metadata": {},
   "source": [
    "## The AST\n",
    "\n",
    "The AST from the above code is fairly straightforward; here is a dump of it:\n",
    "\n",
    "```\n",
    "Module:\n",
    "  Function \n",
    "    Proto 'multiply_transpose' @test/Examples/Toy/Ch1/ast.toy:4:1'\n",
    "    Params: [a, b]\n",
    "    Block {\n",
    "      Return\n",
    "        BinOp: * @test/Examples/Toy/Ch1/ast.toy:5:25\n",
    "          Call 'transpose' [ @test/Examples/Toy/Ch1/ast.toy:5:10\n",
    "            var: a @test/Examples/Toy/Ch1/ast.toy:5:20\n",
    "          ]\n",
    "          Call 'transpose' [ @test/Examples/Toy/Ch1/ast.toy:5:25\n",
    "            var: b @test/Examples/Toy/Ch1/ast.toy:5:35\n",
    "          ]\n",
    "    } // Block\n",
    "  Function \n",
    "    Proto 'main' @test/Examples/Toy/Ch1/ast.toy:8:1'\n",
    "    Params: []\n",
    "    Block {\n",
    "      VarDecl a<> @test/Examples/Toy/Ch1/ast.toy:11:3\n",
    "        Literal: <2, 3>[ <3>[ 1.000000e+00, 2.000000e+00, 3.000000e+00], <3>[ 4.000000e+00, 5.000000e+00, 6.000000e+00]] @test/Examples/Toy/Ch1/ast.toy:11:11\n",
    "      VarDecl b<2, 3> @test/Examples/Toy/Ch1/ast.toy:15:3\n",
    "        Literal: <6>[ 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00] @test/Examples/Toy/Ch1/ast.toy:15:17\n",
    "      VarDecl c<> @test/Examples/Toy/Ch1/ast.toy:19:3\n",
    "        Call 'multiply_transpose' [ @test/Examples/Toy/Ch1/ast.toy:19:11\n",
    "          var: a @test/Examples/Toy/Ch1/ast.toy:19:30\n",
    "          var: b @test/Examples/Toy/Ch1/ast.toy:19:33\n",
    "        ]\n",
    "      VarDecl d<> @test/Examples/Toy/Ch1/ast.toy:22:3\n",
    "        Call 'multiply_transpose' [ @test/Examples/Toy/Ch1/ast.toy:22:11\n",
    "          var: b @test/Examples/Toy/Ch1/ast.toy:22:30\n",
    "          var: a @test/Examples/Toy/Ch1/ast.toy:22:33\n",
    "        ]\n",
    "      VarDecl e<> @test/Examples/Toy/Ch1/ast.toy:25:3\n",
    "        Call 'multiply_transpose' [ @test/Examples/Toy/Ch1/ast.toy:25:11\n",
    "          var: b @test/Examples/Toy/Ch1/ast.toy:25:30\n",
    "          var: c @test/Examples/Toy/Ch1/ast.toy:25:33\n",
    "        ]\n",
    "      VarDecl f<> @test/Examples/Toy/Ch1/ast.toy:28:3\n",
    "        Call 'multiply_transpose' [ @test/Examples/Toy/Ch1/ast.toy:28:11\n",
    "          Call 'transpose' [ @test/Examples/Toy/Ch1/ast.toy:28:30\n",
    "            var: a @test/Examples/Toy/Ch1/ast.toy:28:40\n",
    "          ]\n",
    "          var: c @test/Examples/Toy/Ch1/ast.toy:28:44\n",
    "        ]\n",
    "    } // Block\n",
    "```\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "aba08538",
   "metadata": {},
   "source": [
    "You can reproduce this result and play with the example in the `examples/` directory; try \n",
    "running the next cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cf9b107d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Module:\n",
      "  Function\n",
      "    Proto 'multiply_transpose' @examples/ast.toy:4:1\n",
      "    Params: [a, b]\n",
      "    Block {\n",
      "      Return\n",
      "        BinOp: * @examples/ast.toy:5:25\n",
      "          Call 'transpose' [ @examples/ast.toy:5:10\n",
      "            var: a @examples/ast.toy:5:20\n",
      "          ]\n",
      "          Call 'transpose' [ @examples/ast.toy:5:25\n",
      "            var: b @examples/ast.toy:5:35\n",
      "          ]\n",
      "    } // Block\n",
      "  Function\n",
      "    Proto 'main' @examples/ast.toy:8:1\n",
      "    Params: []\n",
      "    Block {\n",
      "      VarDecl a<> @examples/ast.toy:11:3\n",
      "        Literal: <2, 3>[ <3>[ 1.000000e+00, 2.000000e+00, 3.000000e+00], <3>[ 4.000000e+00, 5.000000e+00, 6.000000e+00]] @examples/ast.toy:11:11\n",
      "      VarDecl b<2, 3> @examples/ast.toy:15:3\n",
      "        Literal: <6>[ 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00] @examples/ast.toy:15:17\n",
      "      VarDecl c<> @examples/ast.toy:19:3\n",
      "        Call 'multiply_transpose' [ @examples/ast.toy:19:11\n",
      "          var: a @examples/ast.toy:19:30\n",
      "          var: b @examples/ast.toy:19:33\n",
      "        ]\n",
      "      VarDecl d<> @examples/ast.toy:22:3\n",
      "        Call 'multiply_transpose' [ @examples/ast.toy:22:11\n",
      "          var: b @examples/ast.toy:22:30\n",
      "          var: a @examples/ast.toy:22:33\n",
      "        ]\n",
      "      VarDecl e<> @examples/ast.toy:25:3\n",
      "        Call 'multiply_transpose' [ @examples/ast.toy:25:11\n",
      "          var: b @examples/ast.toy:25:30\n",
      "          var: c @examples/ast.toy:25:33\n",
      "        ]\n",
      "      VarDecl f<> @examples/ast.toy:28:3\n",
      "        Call 'multiply_transpose' [ @examples/ast.toy:28:11\n",
      "          Call 'transpose' [ @examples/ast.toy:28:30\n",
      "            var: a @examples/ast.toy:28:40\n",
      "          ]\n",
      "          var: c @examples/ast.toy:28:44\n",
      "        ]\n",
      "    } // Block\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from toy.frontend.parser import ToyParser\n",
    "\n",
    "ast_toy = Path(\"examples/ast.toy\")\n",
    "\n",
    "with open(ast_toy) as f:\n",
    "    parser = ToyParser(ast_toy, f.read())\n",
    "\n",
    "print(parser.parse_module().dump())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "025718ef",
   "metadata": {},
   "source": [
    "The code for the lexer is fairly straightforward; it is all in a single file:\n",
    "`toy/lexer.py`. The parser can be found in `toy/parser.py`; it is a recursive \n",
    "descent parser. If you are not familiar with such a Lexer/Parser, these are very similar \n",
    "to the LLVM Kaleidoscope equivalent that are detailed in the first two chapters of the\n",
    "[LLVM Kaleidoscope Tutorial](https://llvm.org/docs/tutorial/MyFirstLanguageFrontend/LangImpl02.html).\n",
    "\n",
    "The next chapter will demonstrate how to convert this AST into MLIR."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
